import torch
import os
from maddpg.actor_critic import Actor, Critic
from torch.optim.lr_scheduler import StepLR  # 添加这一行
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.manual_seed(10)
import logging
from maddpg.compare import compare_models

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    filename='training.log',  # 日志文件名
    filemode='a'  # 追加模式
)

class MADDPG:
    def __init__(self, args, agent_id,num_action,env):  # 因为不同的agent的obs、act维度可能不一样，所以神经网络不同,需要agent_id来区分
        self.args = args
        self.agent_id = agent_id
        self.train_step = 0

        # create the network
        self.actor_network = Actor(args, agent_id,num_action,env).to(device)
        self.critic_network = Critic(args).to(device)

        # build up the target network
        self.actor_target_network = Actor(args, agent_id,num_action,env).to(device)
        self.critic_target_network = Critic(args).to(device)

        # load the weights into the target networks
        self.actor_target_network.load_state_dict(self.actor_network.state_dict())
        self.critic_target_network.load_state_dict(self.critic_network.state_dict())

        # create the optimizer
        self.actor_optim = torch.optim.AdamW(self.actor_network.parameters(),
                                             lr=self.args.lr_actor,
                                             weight_decay = 1e-4)
        self.critic_optim = torch.optim.Adam(self.critic_network.parameters(), lr=self.args.lr_critic)
        
        # 创建学习率调度器
        self.actor_scheduler = StepLR(self.actor_optim, step_size=self.args.lr_step_size, gamma=self.args.lr_gamma)
        self.critic_scheduler = StepLR(self.critic_optim, step_size=self.args.lr_step_size, gamma=self.args.lr_gamma)

        # create the dict for store the model
        if not os.path.exists(self.args.save_dir):
            os.mkdir(self.args.save_dir)
        # path to save the model
        self.model_path = self.args.save_dir + '/' + self.args.scenario_name
        if not os.path.exists(self.model_path):
            os.mkdir(self.model_path)
        self.model_path = self.model_path + '/' + 'agent_%d' % agent_id
        if not os.path.exists(self.model_path):
            os.mkdir(self.model_path)

        # 加载模型
        if os.path.exists(self.model_path + '/actor_params.pkl'):
            self.actor_network.load_state_dict(torch.load(self.model_path + '/actor_params.pkl'))
            self.critic_network.load_state_dict(torch.load(self.model_path + '/critic_params.pkl'))
            print('Agent {} successfully loaded actor_network: {}'.format(self.agent_id,
                                                                          self.model_path + '/actor_params.pkl'))
            print('Agent {} successfully loaded critic_network: {}'.format(self.agent_id,
                                                                           self.model_path + '/critic_params.pkl'))

    # soft update
    def _soft_update_target_network(self):
        for target_param, param in zip(self.actor_target_network.parameters(), self.actor_network.parameters()):
            target_param.data.copy_((1 - self.args.tau) * target_param.data + self.args.tau * param.data)

        for target_param, param in zip(self.critic_target_network.parameters(), self.critic_network.parameters()):
            target_param.data.copy_((1 - self.args.tau) * target_param.data + self.args.tau * param.data)

    # update the network
    def train(self, transitions, other_agents, uav,episode,eval):
        for key in transitions.keys():
            transitions[key] = torch.tensor(transitions[key], dtype=torch.float32,device=device)
        
        # 只使用自身产生的数据
        r = transitions['r_%d' % self.agent_id]  # 训练时只需要自己的reward
        o, u, o_next = [], [], []  # 用来装每个agent经验中的各项
        for agent_id in range(self.args.num_uavs):
            o.append(transitions['o_%d' % agent_id])
            u.append(transitions['u_%d' % agent_id])
            o_next.append(transitions['o_next_%d' % agent_id])

        # calculate the target Q value function
        u_next = []

        # calculate the target Q value function
        with torch.no_grad():
            # 得到下一个状态对应的动作
            index = 0
            for agent_id in range(self.args.num_uavs):
                if agent_id == self.agent_id:
                    u_next.append(self.actor_target_network(o_next[agent_id],uav,episode).to(device))
                else:
                    # 因为传入的other_agents要比总数少一个，可能中间某个agent是当前agent，不能遍历去选择动作
                    u_next.append(other_agents[index].policy.actor_target_network(o_next[agent_id],uav,episode).to(device))
                    index += 1
            q_next = self.critic_target_network(o_next, u_next).detach().to(device)

            target_q = (r.unsqueeze(1) + self.args.gamma * q_next).detach()

        # the q loss
        q_value = self.critic_network(o, u).to(device)
        critic_loss = (target_q - q_value).pow(2).mean()
        l2_loss = torch.tensor(0.).to(device)
        for name, param in self.critic_network.named_parameters():
            if 'weight' in name:
                l2_loss += torch.norm(param)
        critic_loss += self.args.l2_critic * l2_loss
    
        # 使用 f-string 提高性能和可读性
        logging.info(f"Episode {episode}|Train_step {self.train_step+1:03d}| Critic Loss{uav.id}: {critic_loss:.4f}")

        # 获取决策网络第一个线性层的权重
        pre_update = self.actor_network.decision_net[1].weight.data.clone()
            
        # the actor loss
        # 重新选择联合动作中当前agent的动作，其他agent的动作不变
        u[self.agent_id] = self.actor_network(o[self.agent_id],uav,episode).to(device)
        actor_loss = - self.critic_network(o, u).mean().to(device)
        # update the network
        self.actor_optim.zero_grad()
        actor_loss.backward()

        self.actor_optim.step()

        # 在参数更新后添加变化检查
        post_update = self.actor_network.decision_net[1].weight.data.clone()
        diff = torch.norm(post_update - pre_update).item()

        self.critic_optim.zero_grad()
        critic_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.critic_network.parameters(), 1.0)  # 梯度裁剪
        self.critic_optim.step()
        
        # 更保守的目标网络更新
        for target_param, param in zip(self.critic_target_network.parameters(), self.critic_network.parameters()):
            target_param.data.copy_(0.001 * param.data + 0.999 * target_param.data)  # 降低更新速率

        self._soft_update_target_network()
        if self.train_step > 0 and self.train_step % self.args.save_rate == 0 and not eval:
            self.save_model(self.train_step)
            
        # 更新学习率
        if not eval:
            if self.train_step%2 == 0:
                self.actor_scheduler.step()
            self.critic_scheduler.step()
        if eval :
            self.actor_optim = torch.optim.AdamW(self.actor_network.parameters(),
                                             lr=1e-6)
            self.critic_optim = torch.optim.Adam(self.critic_network.parameters(), lr=1e-6)
        self.train_step += 1
        
        return actor_loss.item(), critic_loss.item()

    def save_model(self, train_step):
        num = str(train_step // self.args.save_rate)
        model_path = os.path.join(self.args.save_dir, self.args.scenario_name)
        if not os.path.exists(model_path):
            os.makedirs(model_path)
        model_path = os.path.join(model_path, 'agent_%d' % self.agent_id)
        if not os.path.exists(model_path):
            os.makedirs(model_path)
        torch.save(self.actor_network.state_dict(), model_path + '/' + num + '_actor_params.pkl')
        torch.save(self.critic_network.state_dict(),  model_path + '/' + num + '_critic_params.pkl')
        if int(num) > 1:
            compare_models(model_path + '/' + num + '_actor_params.pkl', model_path + '/' + str(int(num) - 1) + '_actor_params.pkl')


